# import
import numpy as np
import pandas as pd

import time
import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from util.models import *
from util.misc import *
from completeness import *

## mat file
from scipy.io import loadmat

if __name__ == '__main__':
    # load dataset
    # dir_name = os.path.dirname(__file__)
    # matdata = loadmat(dir_name + '/Bruhin_et_al_2010.mat')
    base = os.path.dirname(os.path.abspath(__file__))
    file_name = os.path.normpath(os.path.join(base, '../data/Bruhin_et_al_2010.mat'))
    matdata = loadmat(file_name)

    df_full = pd.DataFrame()
    for key in matdata.keys():
        if not '__' in key:
            df_full[key] = matdata[key].flatten()

    ## drop lotteries over losses
    df = df_full[df_full['z2'] >= 0]

    # preprocessing
    X, Y, subject_idx, cov_idx, cov_list = preprocess(df)

    # main
    # start_time = time.perf_counter()

    alpha_grid_size = 0.01
    eta_grid_size = 0.01
    alpha_ubd = 1
    eta_lbd = -1
    eta_ubd = 5

    ## alpha, eta
    alpha_grid = np.arange(alpha_grid_size, alpha_ubd+alpha_grid_size, alpha_grid_size)
    eta_grid = np.arange(eta_lbd, eta_ubd+eta_grid_size, eta_grid_size)
    param_grid = cartesian_product(alpha_grid, eta_grid)

    completeness =\
        bootstrap_one_iter(X, Y, subject_idx, cov_idx, cov_list, param_grid,
                           base_param=np.array([1,0]), pred=pred_DA, train_model=train_DA)
    print(completeness)